package com.northpool.operator.statistics.dataset.math;

import com.alibaba.fastjson.JSONObject;
import com.northpool.operator.annotation.OperatorInfo;
import com.northpool.operator.common.algorithm.annotation.PluginInfo;
import com.northpool.operator.param.ParamTool;
import com.northpool.operator.type.TypeEnum;
import com.northpool.resources.command.QueryFilter;
import com.northpool.resources.datatable.FieldEncoder;
import com.northpool.resources.datatable.ITable;
import com.northpool.operator.annotation.InputParam;
import com.northpool.operator.annotation.InputParams;
import com.northpool.operator.annotation.OutputParam;
import com.northpool.operator.annotation.OutputParams;
import org.apache.commons.math3.stat.regression.OLSMultipleLinearRegression;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@OperatorInfo(
        ref = "StatisticalMultipleLinearRegression",
        name = "多元线性回归分析"
)
@OutputParams({
        @OutputParam(ref = "statisticalResults", name = "统计结果", type = TypeEnum.STRING)
})
public class StatisticalMultipleLinearRegression  extends BaseMathStatistical {

    @InputParams({
            @InputParam(ref = "dataSourceId", name = "数据源ID", desc = "数据源ID", type = TypeEnum.DATASOURCE, required = true),
            @InputParam(ref = "tableName", name = "数据表名称", desc = "数据表名称，如: public.table", type = TypeEnum.STRING),
            @InputParam(ref = "infoFieldName", name = "字段名称", desc = "字段名称，如: is_passing", type = TypeEnum.STRING),
            @InputParam(ref = "valueFieldName", name = "值字段名称", desc = "值字段名称，如: chinese_score,math_score,physics_score", type = TypeEnum.STRING)
    })
    public StatisticalMultipleLinearRegression() {
    }

    @Override
    protected void executeStatistical(ParamTool paramTool)  {
        String infoFieldName= paramTool.getStr("infoFieldName");

        String valueFieldName= paramTool.getStr("valueFieldName");
        List<String> listValueFields= Arrays.stream(valueFieldName.split(",")).collect(Collectors.toList());
        listValueFields.add(0,infoFieldName);

        try {
            String[] fields = listValueFields.stream().map(a->a).toArray(String[]::new);
            ITable table = tableOperator.getTable(fields, FieldEncoder.ORIGIN_FIELD_ENCODER);

            OLSMultipleLinearRegression regression = new OLSMultipleLinearRegression();
            QueryFilter filter = new QueryFilter();
            List<Map<String, Object>> queryResultData = table.mapDao().query(filter);

            List<Double> listInfo=new ArrayList<>();
            List<Double[]> listValues=new ArrayList<>();

            for (int i = 0; i < queryResultData.size(); i++) {
                Map<String, Object> stringObjectMap = queryResultData.get(i);
                Double infoFieldValue= Double.parseDouble(stringObjectMap.get(infoFieldName).toString()) ;
                listInfo.add(infoFieldValue);

                Double[] valueArr=new Double[listValueFields.size()-1];
                for (int fieldIndex = 1; fieldIndex < listValueFields.size(); fieldIndex++) {
                    valueArr[fieldIndex-1]= Double.parseDouble(stringObjectMap.get(listValueFields.get(fieldIndex)).toString());
                }
                listValues.add(valueArr);
            }

//            this.progress(99);
            double[][] arrayValue = convertListToArray2(listValues);
            double[] arrayInfo = convertListToArray(listInfo);
            regression.newSampleData( arrayInfo,arrayValue);

            // 获取斜率和截距
            double[] coefficients = regression.estimateRegressionParameters();
            double intercept = coefficients[0]; // 截距

            JSONObject jsonObject = new JSONObject();
            jsonObject.put("intercept", intercept);
            jsonObject.put("coefficients", coefficients);

            this.writeResult(jsonObject.toJSONString());

            this.successExit();
        }catch (Exception ex) {
            this.error(ex.getMessage());
            this.failedExit();
        }
    }

    public static double[] convertListToArray(List<Double> list) {
        double[] array = new double[list.size()];
        for (int i = 0; i < list.size(); i++) {
            array[i]=list.get(i);
        }
        return array;
    }

    public static double[][] convertListToArray2(List<Double[]> list) {
        double[][] array = new double[list.size()][];
        for (int i = 0; i < list.size(); i++) {
            Double[] doubleArray = list.get(i);
            array[i] = new double[doubleArray.length];
            for (int j = 0; j < doubleArray.length; j++) {
                array[i][j] = doubleArray[j];
            }
        }
        return array;
    }
}
